{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/anthropic_prompt_caching.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "# Anthropic Prompt Caching\n",
    "\n",
    "In this Notebook, we will demonstrate the usage of [Anthropic Prompt Caching](https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching) with LlamaIndex abstractions.\n",
    "\n",
    "Prompt Caching is enabled by marking `cache_control` in the messages request.\n",
    "\n",
    "\n",
    "## How Prompt Caching works\n",
    "\n",
    "When you send a request with Prompt Caching enabled:\n",
    "\n",
    "1. The system checks if the prompt prefix is already cached from a recent query.\n",
    "2. If found, it uses the cached version, reducing processing time and costs.\n",
    "3. Otherwise, it processes the full prompt and caches the prefix for future use.\n",
    "\n",
    "\n",
    "**Note:**\n",
    "\n",
    "A. Prompt caching works with `Claude 4 Opus`, `Claude 4 Sonnet`, `Claude 3.7 Sonnet`, `Claude 3.5 Sonnet`, `Claude 3.5 Haiku`, `Claude 3 Haiku` and `Claude 3 Opus` models.\n",
    "\n",
    "B. The minimum cacheable prompt length is:\n",
    "\n",
    "    1. 2048 tokens for Claude 3.5 Haiku and Claude 3 Haiku\n",
    "    2. 1024 for all the other models.\n",
    "\n",
    "C. Shorter prompts cannot be cached, even if marked with `cache_control`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup API Keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\n",
    "    \"ANTHROPIC_API_KEY\"\n",
    "] = \"sk-ant-...\"  # replace with your Anthropic API key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.anthropic import Anthropic\n",
    "\n",
    "llm = Anthropic(model=\"claude-3-5-sonnet-20240620\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download Data\n",
    "\n",
    "In this demonstration, we will use the text from the `Paul Graham Essay`. We will cache the text and run some queries based on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-12-14 18:39:03--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.111.133, 185.199.109.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘./paul_graham_essay.txt’\n",
      "\n",
      "./paul_graham_essay 100%[===================>]  73.28K  --.-KB/s    in 0.04s   \n",
      "\n",
      "2024-12-14 18:39:03 (1.62 MB/s) - ‘./paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O './paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[\"./paul_graham_essay.txt\"],\n",
    ").load_data()\n",
    "\n",
    "document_text = documents[0].text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prompt Caching\n",
    "\n",
    "To enable prompt caching, you can just use the `CachePoint` block within LlamaIndex: everything previous to that block will be cached.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can verify if the text is cached by checking the following parameters:\n",
    "\n",
    "`cache_creation_input_tokens:` Number of tokens written to the cache when creating a new entry.\n",
    "\n",
    "`cache_read_input_tokens:` Number of tokens retrieved from the cache for this request.\n",
    "\n",
    "`input_tokens:` Number of input tokens which were not read from or used to create a cache."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import (\n",
    "    ChatMessage,\n",
    "    TextBlock,\n",
    "    CachePoint,\n",
    "    CacheControl,\n",
    ")\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are helpful AI Assitant.\"),\n",
    "    ChatMessage(\n",
    "        role=\"user\",\n",
    "        content=[\n",
    "            TextBlock(\n",
    "                text=f\"{document_text}\",\n",
    "                type=\"text\",\n",
    "            ),\n",
    "            TextBlock(\n",
    "                text=\"\\n\\nWhy did Paul Graham start YC?\",\n",
    "                type=\"text\",\n",
    "            ),\n",
    "            CachePoint(cache_control=CacheControl(type=\"ephemeral\")),\n",
    "        ],\n",
    "    ),\n",
    "]\n",
    "\n",
    "resp = llm.chat(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's examine the raw response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'msg_01PAaZDTjEqcZksFiiqYH42t',\n",
       " 'content': [TextBlock(text='Based on the essay, it seems Paul Graham started Y Combinator (YC) for a few key reasons:\\n\\n1. He had experience as a startup founder with Viaweb and wanted to help other founders avoid mistakes he had made.\\n\\n2. He had ideas about how venture capital could be improved, like making more smaller investments in younger technical founders.\\n\\n3. He was looking for something new to work on after selling Viaweb to Yahoo and trying painting for a while.\\n\\n4. He wanted to gain experience as an investor and thought funding a batch of startups at once would be a good way to do that.\\n\\n5. It started as a \"Summer Founders Program\" to give undergrads an alternative to summer internships, but quickly grew into something more serious.\\n\\n6. He saw an opportunity to scale startup funding by investing in batches of companies at once.\\n\\n7. He was excited by the potential to help create new startups and technologies.\\n\\n8. It allowed him to continue working with his friends/former colleagues Robert Morris and Trevor Blackwell.\\n\\n9. He had built an audience through his essays that provided deal flow for potential investments.\\n\\nSo in summary, it was a combination of wanting to help founders, improve venture capital, gain investing experience, work with friends, and leverage his existing audience/expertise in the startup world. The initial idea evolved quickly from a summer program into a new model for seed investing.', type='text')],\n",
       " 'model': 'claude-3-5-sonnet-20240620',\n",
       " 'role': 'assistant',\n",
       " 'stop_reason': 'end_turn',\n",
       " 'stop_sequence': None,\n",
       " 'type': 'message',\n",
       " 'usage': Usage(input_tokens=4, output_tokens=305, cache_creation_input_tokens=9, cache_read_input_tokens=17467)}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resp.raw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, since I've ran this a few different times, `cache_creation_input_tokens` and `cache_read_input_tokens` are both higher than zero, indicating that the text was cached properly.\n",
    "\n",
    "Now, let’s run another query on the same document. It should retrieve the document text from the cache, which will be reflected in `cache_read_input_tokens`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are helpful AI Assitant.\"),\n",
    "    ChatMessage(\n",
    "        role=\"user\",\n",
    "        content=[\n",
    "            TextBlock(\n",
    "                text=f\"{document_text}\",\n",
    "                type=\"text\",\n",
    "            ),\n",
    "            TextBlock(\n",
    "                text=\"\\n\\nWhat did Paul Graham do growing up?\",\n",
    "                type=\"text\",\n",
    "            ),\n",
    "            CachePoint(cache_control=CacheControl(type=\"ephemeral\")),\n",
    "        ],\n",
    "    ),\n",
    "]\n",
    "\n",
    "resp = llm.chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'msg_011TQgbpBuBkZAJeatVVcqtp',\n",
       " 'content': [TextBlock(text='Based on the essay, here are some key things Paul Graham did growing up:\\n\\n1. As a teenager, he focused mainly on writing and programming outside of school. He tried writing short stories but says they were \"awful\".\\n\\n2. At age 13-14, he started programming on an IBM 1401 computer at his school district\\'s data processing center. He used an early version of Fortran.\\n\\n3. In high school, he convinced his father to buy a TRS-80 microcomputer around 1980. He wrote simple games, a program to predict model rocket flight, and a word processor his father used.\\n\\n4. He went to college intending to study philosophy, but found it boring. He then decided to switch to studying artificial intelligence (AI).\\n\\n5. In college, he learned Lisp programming language, which expanded his concept of what programming could be. \\n\\n6. For his undergraduate thesis, he reverse-engineered SHRDLU, an early natural language processing program.\\n\\n7. He applied to grad schools for AI and ended up going to Harvard for graduate studies.\\n\\n8. In grad school, he realized AI as practiced then was not going to achieve true intelligence. He pivoted to focusing more on Lisp programming.\\n\\n9. He started writing a book about Lisp hacking while in grad school, which was eventually published in 1993 as \"On Lisp\".\\n\\nSo in summary, his early years were focused on writing, programming (especially Lisp), and studying AI, before he eventually moved on to other pursuits after grad school. The essay provides a detailed account of his intellectual development in these areas.', type='text')],\n",
       " 'model': 'claude-3-5-sonnet-20240620',\n",
       " 'role': 'assistant',\n",
       " 'stop_reason': 'end_turn',\n",
       " 'stop_sequence': None,\n",
       " 'type': 'message',\n",
       " 'usage': Usage(input_tokens=4, output_tokens=356, cache_creation_input_tokens=0, cache_read_input_tokens=17476)}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resp.raw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the response was generated using cached text, as indicated by `cache_read_input_tokens`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With Anthropic, the default cache lasts 5 minutes. You can also have longer lasting caches, for instance 1 hour, you just have to specify that under the `ttl` argument in `CachControl`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are helpful AI Assitant.\"),\n",
    "    ChatMessage(\n",
    "        role=\"user\",\n",
    "        content=[\n",
    "            TextBlock(\n",
    "                text=f\"{document_text}\",\n",
    "                type=\"text\",\n",
    "            ),\n",
    "            TextBlock(\n",
    "                text=\"\\n\\nWhat did Paul Graham do growing up?\",\n",
    "                type=\"text\",\n",
    "            ),\n",
    "            CachePoint(\n",
    "                cache_control=CacheControl(type=\"ephemeral\", ttl=\"1h\"),\n",
    "            ),\n",
    "        ],\n",
    "    ),\n",
    "]\n",
    "\n",
    "resp = llm.chat(messages)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
